import os
import traceback

import pandas as pd
from datetime import timedelta

from base.constant import out_folder
from base.util.file_util import check_path
from set_label.imperva import get_dangerous_from_imperva
from set_label.splunk import get_dangerous_ips


def set_label(config_yaml):
    print("[打标签]开始，结束后记得启动 arkime capture")
    allTrue = config_yaml['label']['allTrue']
    # 处理原始数据
    duration = config_yaml['task_package']['duration']
    session_list = pd.read_csv(f"{out_folder}{os.sep}{config_yaml['session']['time_data_out']}", delimiter=',')
    session_list = session_list.fillna('')
    # useless
    session_list['same_src_dst_size_mean'] = 0
    session_list['same_src_dst_size_var'] = 0
    session_list['packet_size_mean'] = pd.to_numeric(session_list['packet_size_mean'], errors='coerce')
    session_list['same_src_dst_size_mean'] = session_list.groupby(['destination.ip', 'destination.port'])[
        'packet_size_mean'].transform('mean')
    session_list['same_src_dst_size_var'] = session_list.groupby(['destination.ip', 'destination.port'])[
        'packet_size_mean'].transform('var')
    session_list['source.ip'] = session_list['source.ip'].astype(str)
    session_list['source.port'] = session_list['source.port'].astype(str)
    session_list['destination.ip'] = session_list['destination.ip'].astype(str)
    session_list['destination.port'] = session_list['destination.port'].astype(str)
    # firstPacket和lastPacket是13位数的时间戳
    session_list['firstPacket'] = session_list['firstPacket'].astype(int)
    session_list['lastPacket'] = session_list['lastPacket'].astype(int)
    # 增加时间列用于判断
    session_list['firstPacket_datetime'] = pd.to_datetime(session_list['firstPacket'],
                                                          unit='ms').dt.tz_localize('UTC')
    session_list['lastPacket_datetime'] = pd.to_datetime(session_list['lastPacket'],
                                                         unit='ms').dt.tz_localize('UTC')
    session_list['firstPacket_datetime'] = session_list['firstPacket_datetime'].dt.tz_convert('Asia/Shanghai')
    session_list['lastPacket_datetime'] = session_list['lastPacket_datetime'].dt.tz_convert('Asia/Shanghai')
    # 增加列
    session_list['isDangerous'] = False
    # 后面匹配到异常字段则填充
    session_list['PROTOCOL'] = ''
    session_list['DENY_METHOD'] = ''
    session_list['THREAT_SUMMARY'] = ''
    session_list['SEVERITY'] = ''
    if allTrue:
        session_list.drop(columns=['firstPacket_datetime', 'lastPacket_datetime'], inplace=True)
        session_list['isDangerous'] = True
        session_list.to_csv(check_path(f"{out_folder}{os.sep}{config_yaml['label']['label_out']}"), mode='w',
                            index=False)
        print("[打标签]打标签结束")
    else:
        start_time = config_yaml['session']['start_time']
        end_time = config_yaml['session']['end_time']
        dangerous_ids = get_dangerous_ips(config_yaml, "")
        dangerous_waf = get_dangerous_from_imperva(start_time, end_time, get_new=False)
        dangerous_waf['type'] = 'waf'
        dangerous_ids['type'] = 'ids'
        dangerous_all = pd.concat([dangerous_ids, dangerous_waf])
        if len(dangerous_all) == 0:
            print("[打标签]异常数量为空，打标签结束")
            session_list.drop(columns=['firstPacket_datetime', 'lastPacket_datetime'], inplace=True)

        else:
            # 处理异常数据
            dangerous_all.fillna('', inplace=True)
            dangerous_all['SIP'] = dangerous_all['SIP'].astype(str)
            dangerous_all['S_PORT'] = dangerous_all['S_PORT'].astype(str)
            dangerous_all['DIP'] = dangerous_all['DIP'].astype(str)
            dangerous_all['D_PORT'] = dangerous_all['D_PORT'].astype(str)
            dangerous_all['THREAT_TIME'] = pd.to_datetime(dangerous_all['THREAT_TIME'], format='%Y-%m-%d %H:%M:%S')
            try:
                count = 0
                unmatched_data = []
                for ind, (i, d_item) in enumerate(dangerous_all.iterrows()):
                    type = d_item['type']
                    if type == 'ids':
                        sip = session_list['source.ip']
                    else:
                        sip = session_list['http.xffIp'].apply(
                            lambda x: x[0] if isinstance(x, list) and len(x) > 0 else "")
                    threat_time = d_item['THREAT_TIME'].tz_localize('Asia/Shanghai')
                    threat_start_time = threat_time - timedelta(minutes=duration)
                    threat_end_time = threat_time + timedelta(minutes=duration)
                    result_session = session_list[(
                            (sip == d_item["SIP"]) &
                            (session_list['source.port'] == d_item["S_PORT"]) &
                            (session_list['destination.ip'] == d_item["DIP"]) &
                            (session_list['destination.port'] == d_item["D_PORT"])
                            & (session_list['firstPacket_datetime'] >= threat_start_time)
                            & (session_list['lastPacket_datetime'] <= threat_end_time)
                    )]
                    count = count + len(result_session)
                    if len(result_session) != 0:
                        print(
                            f'[打标签]进度：{ind + 1}/{len(dangerous_all)}，本次循环匹配到：{len(result_session)},共匹配到：{count}')
                        # if len(result_session) > 1:
                        #     result_session = result_session[result_session['http.xffIp'] == d_item["XFF_IP"]]
                        #     if len(result_session) > 0:
                        #         result_session = result_session
                        session_list.loc[session_list['id'].isin(result_session['id']), [
                            'DENY_METHOD', 'PROTOCOL', 'SEVERITY', 'THREAT_SUMMARY', 'isDangerous']] = [
                            d_item['DENY_METHOD'],
                            d_item['PROTOCOL'],
                            d_item['SEVERITY'],
                            d_item['THREAT_SUMMARY'],
                            True
                        ]
                    else:
                        if config_yaml['label']['is_output_unmatch']:
                            unmatched_data.append(d_item)
                            print(
                                f"[打标签]进度：{ind + 1}/{len(dangerous_all)}｜没有匹配到：{d_item['SIP']}:{d_item['S_PORT']}---{d_item['DIP']}:{d_item['D_PORT']}")
                if config_yaml['label']['is_output_unmatch']:
                    if len(unmatched_data) > 0:
                        print(f"[打标签]完成，有{len(unmatched_data)}条没有匹配到")
                        if not allTrue:
                            pd.DataFrame(unmatched_data).to_csv(
                                check_path(f"{out_folder}{os.sep}{config_yaml['label']['label_unmatch_out']}"), mode='w',
                                index=False, errors="replace")
                session_list.drop(columns=['firstPacket_datetime', 'lastPacket_datetime'], inplace=True)
                if allTrue:
                    session_list['isDangerous'] = True
                session_list.to_csv(check_path(f"{out_folder}{os.sep}{config_yaml['label']['label_out']}"), mode='w',
                                    index=False, errors="replace")
            except Exception as e:
                print("[打标签]报错：")
                print(e)
                traceback.print_exc()
                exit(0)
